import cv2
import numpy as np
import matplotlib.pyplot as plt

def extract_watermark(carrier_path, embedded_path, alpha=0.5, output_path="extracted_watermark.png"):
    # 读取载体图像和嵌入水印图像
    carrier = cv2.imread(carrier_path, cv2.IMREAD_GRAYSCALE)
    embedded = cv2.imread(embedded_path, cv2.IMREAD_GRAYSCALE)

    # 计算DFT
    carrier_dft = np.fft.fftshift(np.fft.fft2(carrier))
    embedded_dft = np.fft.fftshift(np.fft.fft2(embedded))

    # 提取水印 (频域操作)
    watermark_dft = (embedded_dft - carrier_dft) / alpha

    # 逆DFT还原水印图像
    watermark_idft_shifted = np.fft.ifftshift(watermark_dft)
    watermark_image = np.abs(np.fft.ifft2(watermark_idft_shifted))

    # 保存结果
    plt.imsave(output_path, watermark_image, cmap='gray')
    print(f"提取的水印图像已保存到 {output_path}")

# 调用
extract_watermark("alena.bmp", "embedded_image.png")

